/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file flash_attention_score_grad_tiling_bng_gs1s2.cc
 * \brief
 */

#include "flash_attention_score_grad_tiling_bng_gs1s2.h"
#include "tiling/tiling_type.h"
#include "tiling/tiling_templates_registry.h"

using namespace AscendC;
using namespace ge;

namespace optiling {

constexpr uint32_t BYTE_BLOCK = 32;
// 分形大小
constexpr uint32_t FRACTAL_NUM = 16;

constexpr uint32_t gScaleAttrIdx = 0;
constexpr uint32_t gKeepProbAttrIdx = 1;
constexpr uint32_t gPrevTokenAttrIdx = 2;
constexpr uint32_t gNextTokenAttrIdx = 3;
constexpr uint32_t gHeadNumAttrIdx = 4;
constexpr uint32_t gLayoutAttrIdx = 5;

#define KB 1024
#define UB_SIZE (192 * KB)
#define L1_SIZE (512 * KB)
#define L0AB_SIZE (64 * KB)
#define L0C_SIZE (128 * KB)

uint32_t CeilFuncV2(uint32_t num1, uint32_t num2) {
  if (num2 == 0) {
    return 0;
  }
  return (num1 + num2 - 1) / num2;
}

bool FlashAttentionScoreGradTilingBNGgs1s2::IsCapable() {
  return false; //zyt temp
}

uint64_t FlashAttentionScoreGradTilingBNGgs1s2::GetTilingKey() const {
  auto queryType = context_->GetInputDesc(0)->GetDataType();
  uint64_t tilingKey;
  // current version is always true
  bool isHighPrecision = true;
  if (queryType == ge::DT_FLOAT16) {
    if (isHighPrecision) {
      // 3012
      tilingKey = TILINGKEY(G, N2, B, FLOAT16_PRECISION, BSND, ALL);
    } else {
      // 0012
      tilingKey = TILINGKEY(G, N2, B, FLOAT16, BSND, ALL);
    }
  } else if (queryType == ge::DT_BF16) {
    // 2012
    tilingKey = TILINGKEY(G, N2, B, BFLOAT16, BSND, ALL);
  } else {
    // 1012
    tilingKey = TILINGKEY(G, N2, B, FLOAT32, BSND, ALL);
  }
  return tilingKey;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::GetPlatformInfo() {
  // 1. Get CoreNum
  auto platformInfoPtr = context_->GetPlatformInfo();
  if (platformInfoPtr == nullptr) {
    auto compileInfoPtr = reinterpret_cast<const FlashAttentionScoreGradCompileInfo*>(context_->GetCompileInfo());
    OPS_ERR_IF(compileInfoPtr == nullptr, OPS_REPORT_CUBE_INNER_ERR(context_, "compile_info is null"),
               return ge::GRAPH_FAILED);

    mBaseInfoParams.coreNum = compileInfoPtr->aivNum;
    mBaseInfoParams.aicNum = compileInfoPtr->aicNum;
  } else {
    ascendcPlatform_ = std::unique_ptr<platform_ascendc::PlatformAscendC>(
        new (std::nothrow) platform_ascendc::PlatformAscendC(platformInfoPtr));
    mBaseInfoParams.coreNum = ascendcPlatform_->GetCoreNumAiv();
    mBaseInfoParams.aicNum = ascendcPlatform_->GetCoreNumAic();
  }

  // 2. Get buffer size
  mBaseInfoParams.ubSize = UB_SIZE;
  mBaseInfoParams.l1Size = L1_SIZE;
  mBaseInfoParams.l0aSize = L0AB_SIZE;
  mBaseInfoParams.l0bSize = L0AB_SIZE;
  mBaseInfoParams.l0cSize = L0C_SIZE;

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::GetShapeAttrsInfo() {
  uint32_t b = 0;
  uint32_t n1 = 0;
  uint32_t n2 = 0;
  uint32_t d = 0;
  uint32_t h = 0;
  uint32_t s1 = 0;
  uint32_t s2 = 0;
  uint32_t g = 1;
  uint32_t layoutType = BSH;

  const char* inputLayout = context_->GetAttrs()->GetAttrPointer<char>(gLayoutAttrIdx);
  const gert::StorageShape* queryShape = context_->GetInputShape(0);
  const gert::StorageShape* keyShape = context_->GetInputShape(1);

  b = queryShape->GetStorageShape().GetDim(DIM_INDEX_0);
  s1 = queryShape->GetStorageShape().GetDim(DIM_INDEX_1);
  h = queryShape->GetStorageShape().GetDim(DIM_INDEX_2);
  n1 = *context_->GetAttrs()->GetAttrPointer<int>(gHeadNumAttrIdx);
  n2 = n1;
  if (inputLayout[LAYOUT_INDEX_0] == 'B' && inputLayout[LAYOUT_INDEX_1] == 'S' && inputLayout[LAYOUT_INDEX_2] == 'H') {
    s2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_1);
    d = h / n1;
  } else if (inputLayout[LAYOUT_INDEX_0] == 'S' && inputLayout[LAYOUT_INDEX_1] == 'B' &&
             inputLayout[LAYOUT_INDEX_2] == 'H') {
    b = queryShape->GetStorageShape().GetDim(DIM_INDEX_1);
    s1 = queryShape->GetStorageShape().GetDim(DIM_INDEX_0);
    s2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_0);
    d = h / n1;
    layoutType = SBH;
  } else if (inputLayout[LAYOUT_INDEX_0] == 'B' && inputLayout[LAYOUT_INDEX_1] == 'N' &&
             inputLayout[LAYOUT_INDEX_2] == 'S' && inputLayout[LAYOUT_INDEX_3] == 'D') {
    n1 = queryShape->GetStorageShape().GetDim(DIM_INDEX_1);
    s1 = queryShape->GetStorageShape().GetDim(DIM_INDEX_2);
    n2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_1);
    s2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_2);
    d = queryShape->GetStorageShape().GetDim(DIM_INDEX_3);
    g = n1 / n2;
    layoutType = BNSD;
  } else if (inputLayout[LAYOUT_INDEX_0] == 'B' && inputLayout[LAYOUT_INDEX_1] == 'S' &&
             inputLayout[LAYOUT_INDEX_2] == 'N' && inputLayout[LAYOUT_INDEX_3] == 'D') {
    n1 = queryShape->GetStorageShape().GetDim(DIM_INDEX_2);
    s2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_1);
    n2 = keyShape->GetStorageShape().GetDim(DIM_INDEX_2);
    d = queryShape->GetStorageShape().GetDim(DIM_INDEX_3);
    g = n1 / n2;
    layoutType = BSND;
  }

  mBaseInfoParams.b = b;
  mBaseInfoParams.n1 = n1;
  mBaseInfoParams.n2 = n2;
  mBaseInfoParams.s1 = s1;
  mBaseInfoParams.s2 = s2;
  mBaseInfoParams.d = d;
  mBaseInfoParams.g = g;
  mBaseInfoParams.h = h;
  mBaseInfoParams.layout = layoutType;
  // sparse
  int preTokens = *context_->GetAttrs()->GetAttrPointer<int>(gPrevTokenAttrIdx);
  int nextTokens = *context_->GetAttrs()->GetAttrPointer<int>(gNextTokenAttrIdx);
  mBaseInfoParams.isSparse =
      ((uint32_t)preTokens < mBaseInfoParams.s1 || (uint32_t)nextTokens < mBaseInfoParams.s1) ? true : false;

  // data type
  auto queryType = context_->GetInputTensor(0)->GetDataType();
  if (queryType == ge::DT_FLOAT16) {
    mBaseInfoParams.dataTypeSize = DATA_SIZE_FP16;
    mBaseInfoParams.mask = VMASK_FP16;
    mBaseInfoParams.dataType = 1;
  } else if (queryType == ge::DT_BF16) {
    mBaseInfoParams.dataTypeSize = DATA_SIZE_FP16;
    mBaseInfoParams.mask = VMASK_FP16;
    mBaseInfoParams.dataType = 2;
  } else {
    mBaseInfoParams.dataTypeSize = DATA_SIZE_FP32;
    mBaseInfoParams.mask = VMASK_FP32;
    mBaseInfoParams.dataType = 0;
  }

  mBaseInfoParams.calTypeSize = DATA_SIZE_FP32;
  mBaseInfoParams.dType = queryType;
  mBaseInfoParams.blockNum = BYTE_BLOCK / mBaseInfoParams.dataTypeSize;

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::SetBaseTilingData() {
  // base tiling data
  tilingData.gBNBaseParams.set_coreNum(mBaseInfoParams.coreNum);

  tilingData.gBNBaseParams.set_B(mBaseInfoParams.b);
  tilingData.gBNBaseParams.set_N1(mBaseInfoParams.n1);
  tilingData.gBNBaseParams.set_N2(mBaseInfoParams.n2);
  tilingData.gBNBaseParams.set_S1(mBaseInfoParams.s1);
  tilingData.gBNBaseParams.set_S2(mBaseInfoParams.s2);
  tilingData.gBNBaseParams.set_D(mBaseInfoParams.d);
  tilingData.gBNBaseParams.set_G(mBaseInfoParams.g);
  tilingData.gBNBaseParams.set_H(mBaseInfoParams.h);
  tilingData.gBNBaseParams.set_batch(mBaseInfoParams.b * mBaseInfoParams.n1);
  tilingData.gBNBaseParams.set_layout(mBaseInfoParams.layout);

  tilingData.gBNBaseParams.set_typeBlockNum(BLOCK_SIZE / sizeof(float));
  tilingData.gBNBaseParams.set_scaleValue(*(context_->GetAttrs()->GetAttrPointer<float>(gScaleAttrIdx)));
  tilingData.gBNBaseParams.set_keepProb(*(context_->GetAttrs()->GetAttrPointer<float>(gKeepProbAttrIdx)));

  int preTokens = *context_->GetAttrs()->GetAttrPointer<int>(gPrevTokenAttrIdx);
  int nextTokens = *context_->GetAttrs()->GetAttrPointer<int>(gNextTokenAttrIdx);
  tilingData.gBNBaseParams.set_preTokens(preTokens);
  tilingData.gBNBaseParams.set_nextTokens(nextTokens);

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::SetCoresSplitTilingData() {
  uint32_t gInner = mBaseInfoParams.gInner;
  uint32_t gOut = CeilFuncV2(mBaseInfoParams.g, gInner);
  uint32_t gTail = mBaseInfoParams.g % gInner == 0 ? gInner : mBaseInfoParams.g % gInner;
  uint32_t innerMatResSize = mBaseInfoParams.s1 * mBaseInfoParams.s2 * mBaseInfoParams.calTypeSize;

  tilingData.gBNSplitCoreParams.set_gOut(gOut);
  tilingData.gBNSplitCoreParams.set_gInner(gInner);
  tilingData.gBNSplitCoreParams.set_gTail(gTail);
  tilingData.gBNSplitCoreParams.set_innerMatResSize(innerMatResSize);

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::SetUbSplitTilingData() {
  uint32_t batch = tilingData.gBNBaseParams.get_batch();

  uint32_t batchRange = CeilFuncV2(batch, mBaseInfoParams.coreNum);
  uint32_t usedCoreNum = CeilFuncV2(batch, batchRange);
  uint32_t remainBatchRange = batch - batchRange * (usedCoreNum - 1);

  mBaseInfoParams.coreNum = usedCoreNum;
  tilingData.gBNBaseParams.set_coreNum(mBaseInfoParams.coreNum);

  tilingData.gBNSplitCoreParams.set_batchRange(batchRange);
  tilingData.gBNSplitCoreParams.set_splitedBatchRange(batchRange);
  tilingData.gBNSplitCoreParams.set_splitedRemainBatchRange(remainBatchRange);

  return ge::GRAPH_SUCCESS;
}

// gInner 从 256 逐步缩小, 找到适合UB大小的最大 gInner
uint32_t FlashAttentionScoreGradTilingBNGgs1s2::CalcGInner() {
  uint32_t gInner = std::min(mBaseInfoParams.g, uint32_t(256 - 16));
  // to avoid repeat max value 255
  while (!CheckGInnerLegal(gInner)) {
    if (gInner <= FRACTAL_NUM) {
      return 0;
    }
    gInner = gInner - FRACTAL_NUM;
  }
  return gInner;
};

bool FlashAttentionScoreGradTilingBNGgs1s2::CheckGInnerLegal(uint32_t gInner) {
  uint32_t coexNode = 6;
  // 6: 3 TQue, 3 Tbuff
  // input scene
  uint32_t baseS1S2 = gInner * mBaseInfoParams.s1 * mBaseInfoParams.s2 * mBaseInfoParams.calTypeSize;
  uint32_t baseS1D = gInner * mBaseInfoParams.s1 * mBaseInfoParams.d * mBaseInfoParams.calTypeSize;
  uint32_t baseS2D = gInner * mBaseInfoParams.s2 * mBaseInfoParams.d * mBaseInfoParams.calTypeSize;
  uint32_t maxBufferSize = std::max(baseS1S2, std::max(baseS1D, baseS2D));

  // simplesoftmax and dropout
  auto softmaxShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
  auto dropoutShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
  auto softmaxgradShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
  uint32_t softmaxTmpSize = GetSoftMaxMinTmpSize(softmaxShape, mBaseInfoParams.calTypeSize, true);
  uint32_t dropoutTmpSize = GetDropOutMinTmpSize(dropoutShape, mBaseInfoParams.calTypeSize, true);
  uint32_t softmaxgradTmpSize = GetSoftMaxGradMinTmpSize(softmaxgradShape, mBaseInfoParams.calTypeSize, true, false);
  uint32_t maxTmpBufferSize = std::max(softmaxTmpSize, std::max(dropoutTmpSize, softmaxgradTmpSize));

  maxBufferSize = std::max(maxTmpBufferSize, maxBufferSize);
  uint32_t baseUB = coexNode * maxBufferSize;

  if (baseUB > mBaseInfoParams.ubSize) {
    return false;
  }

  // Loc buffer
  if (baseS1S2 > mBaseInfoParams.l0cSize) {
    return false;
  }

  return true;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::DoOpTiling() {
  SetBaseTilingData();
  SetCoresSplitTilingData();
  SetUbSplitTilingData();

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::DoLibApiTiling() {
  // mm tiling
  matmul_tiling::MatmulApiTiling mm1;
  matmul_tiling::MatmulApiTiling mm2;
  matmul_tiling::MatmulApiTiling mm3;

  // s轴不切，所以默认singleM与originM等大，一个入参
  // BMM1 BMM2
  SetMatmulTiling1(mm1, tilingData.mm1TilingData);

  // BMM4 BMM3.2
  SetMatmulTiling2(mm2, tilingData.mm2TilingData);

  // BMM3.1
  SetMatmulTiling3(mm3, tilingData.mm3TilingData);

  // softmax
  auto softmaxShape = Shape({mBaseInfoParams.s1, mBaseInfoParams.s2});
  SoftMaxTilingFunc(softmaxShape, sizeof(float), (mBaseInfoParams.s1 * mBaseInfoParams.s2),
                    tilingData.softmaxTilingData);

  // softmaxgrad
  SoftMaxGradTilingFunc(softmaxShape, mBaseInfoParams.calTypeSize, (mBaseInfoParams.s1 * mBaseInfoParams.s2),
                        tilingData.softmaxGradTilingData, true);

  return ge::GRAPH_SUCCESS;
}

void FlashAttentionScoreGradTilingBNGgs1s2::SetMatmulTiling1(matmul_tiling::MatmulApiTiling& mm,
                                                             TCubeTiling& mmTiling) {
  // set matmul type
  mm.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16, false);
  mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16, true);
  mm.SetCType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

  // set base shape
  mm.SetShape(mBaseInfoParams.s1, mBaseInfoParams.s2, mBaseInfoParams.d);

  uint32_t alignNBlockNum = (mBaseInfoParams.s2 + mBaseInfoParams.blockNum - 1) / mBaseInfoParams.blockNum;
  uint32_t alignNSize = alignNBlockNum * mBaseInfoParams.blockNum;
  // set origin shape
  if (mBaseInfoParams.layout == BSH || mBaseInfoParams.layout == BSND) {
    // BSH or BSND
    mm.SetOrgShape(mBaseInfoParams.s1, alignNSize, mBaseInfoParams.d * mBaseInfoParams.n2 * mBaseInfoParams.g,
                   mBaseInfoParams.d * mBaseInfoParams.n2);
  } else if (mBaseInfoParams.layout == SBH) {
    // SBH
    mm.SetOrgShape(mBaseInfoParams.s1, alignNSize,
                   mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2 * mBaseInfoParams.g,
                   mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2);
  } else {
    // BNSD
    mm.SetOrgShape(mBaseInfoParams.s1, alignNSize, mBaseInfoParams.d);
  }
  mm.SetBias(false);
  mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * alignNSize);
  // L1, L0C, UB, bt
  mm.SetFixSplit(mBaseInfoParams.s1, mBaseInfoParams.s1, -1);
  mm.GetTiling(mmTiling);

  mmTiling.set_shareMode(0);
  mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
  mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

void FlashAttentionScoreGradTilingBNGgs1s2::SetMatmulTiling2(matmul_tiling::MatmulApiTiling& mm,
                                                             TCubeTiling& mmTiling) {
  // set matmul type
  mm.SetAType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16,
              true);
  mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16, true);
  mm.SetCType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

  // set base shape
  mm.SetShape(mBaseInfoParams.s2, mBaseInfoParams.d, mBaseInfoParams.s1);

  // set origin shape
  if (mBaseInfoParams.layout == BSH || mBaseInfoParams.layout == BSND) {
    // BSH or BSND
    mm.SetOrgShape(mBaseInfoParams.s2, mBaseInfoParams.d * mBaseInfoParams.n2 * mBaseInfoParams.g, mBaseInfoParams.s1,
                   mBaseInfoParams.s1);
  } else if (mBaseInfoParams.layout == SBH) {
    // SBH
    mm.SetOrgShape(mBaseInfoParams.s2, mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2 * mBaseInfoParams.g,
                   mBaseInfoParams.s1, mBaseInfoParams.s1);
  } else {
    // BNSD
    mm.SetOrgShape(mBaseInfoParams.s2, mBaseInfoParams.d, mBaseInfoParams.s1, mBaseInfoParams.s1);
  }
  mm.SetBias(false);
  mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * mBaseInfoParams.s1);
  // L1, L0C, UB, bt
  mm.SetFixSplit(mBaseInfoParams.s2, mBaseInfoParams.d, -1);
  mm.GetTiling(mmTiling);

  mmTiling.set_shareMode(0);
  mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
  mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

void FlashAttentionScoreGradTilingBNGgs1s2::SetMatmulTiling3(matmul_tiling::MatmulApiTiling& mm,
                                                             TCubeTiling& mmTiling) {
  // set matmul type
  mm.SetAType(matmul_tiling::TPosition::VECCALC, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16,
              false);
  mm.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_BFLOAT16, false);
  mm.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);

  // set base shape
  mm.SetShape(mBaseInfoParams.s1, mBaseInfoParams.d, mBaseInfoParams.s2);

  // set origin shape
  if (mBaseInfoParams.layout == BSH || mBaseInfoParams.layout == BSND) {
    // BSH or BSND
    mm.SetOrgShape(mBaseInfoParams.s1, mBaseInfoParams.d * mBaseInfoParams.n2, mBaseInfoParams.s2, mBaseInfoParams.s2);
  } else if (mBaseInfoParams.layout == SBH) {
    // SBH
    mm.SetOrgShape(mBaseInfoParams.s1, mBaseInfoParams.d * mBaseInfoParams.b * mBaseInfoParams.n2, mBaseInfoParams.s2,
                   mBaseInfoParams.s2);
  } else {
    // BNSD
    mm.SetOrgShape(mBaseInfoParams.s1, mBaseInfoParams.d, mBaseInfoParams.s2, mBaseInfoParams.s2);
  }
  mm.SetBias(false);
  mm.SetBufferSpace(-1, -1, mBaseInfoParams.s1 * mBaseInfoParams.s2);
  // L1, L0C, UB, bt
  mm.SetFixSplit(mBaseInfoParams.s1, mBaseInfoParams.d, -1);
  mm.GetTiling(mmTiling);

  mmTiling.set_shareMode(0);
  mmTiling.set_shareL1Size(mBaseInfoParams.l1Size);
  mmTiling.set_shareL0CSize(mBaseInfoParams.l0cSize);
}

ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::GetWorkspaceSize() {
  const uint32_t baseSize = 100 * 1024 * 1024;
  const uint32_t minCoreNum = 4;
  auto aicNum = mBaseInfoParams.aicNum;
  auto aivNum = mBaseInfoParams.coreNum;
  uint32_t coreNum = tilingData.gBNBaseParams.get_coreNum();
  uint32_t innerMatResSize = tilingData.gBNSplitCoreParams.get_innerMatResSize();
  context_->SetBlockDim(CalcTschBlockDim(coreNum, aicNum, aivNum));
  size_t* workspaces = context_->GetWorkspaceSizes(1);
  workspaces[0] = baseSize + innerMatResSize * coreNum * coreNum * minCoreNum;

  return ge::GRAPH_SUCCESS;
}
ge::graphStatus FlashAttentionScoreGradTilingBNGgs1s2::PostTiling() {
  tilingData.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
  context_->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
  return ge::GRAPH_SUCCESS;
}

}  // namespace optiling
